from template import *
from collect import collect

t = 0
dt = 0.5e-12 * 500000 / 1.0e-6   # unit is us
dt = 1


val1 = collect("/Diagnostic/", "particle_number")

nx = (val1.shape)[0]
x = np.linspace(0,nx*dt,nx)

##inite the fig of matplotlib
fig=plt.figure(figsize=(10,10))
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.6,hspace=0.55)


#============Total particle number=======================================
val1 = collect("/Diagnostic/", "particle_number")
val1_1d = val1[:, 0, 0, 0]


val2 = collect("/Diagnostic/", "particle_number")
val2_1d = val2[:, 0, 0, 1]


ax0=fig.add_subplot(5,1,1)


ax0.yaxis.set_major_formatter(yformatter)

line0=ax0.plot(x, val1_1d, label='Electron', linestyle = linestyles[0])
line0=ax0.plot(x, val2_1d, label=r'$\mathrm{D^+}$ ion', linestyle = linestyles[1])

xmin = 0
xmax = x.max()
ymin = 0.0
ymax = 1.2 * max( val1_1d.max(), val2_1d.max() )

ax0.set_xlim((xmin, xmax))
#ax0.set_ylim((ymin, ymax))
ax0.set_ylabel("N", fontsize = label_fontsize)
#ax0.yaxis.set_label_coords(ylabel_x, 0.5)



#ax0.grid(True)
ax0.legend(framealpha=0.5)

ax0.annotate("(a)", xy=get_axis_limits(ax0), annotation_clip=False)

##============Particle flux======================================================
val1 = collect("/Diagnostic/", "particle_flux_left")
val1_1d = val1[:, 0, 0, 0]

val2 = collect("/Diagnostic/", "particle_flux_left")
val2_1d = val2[:, 0, 0, 1]


ax0=fig.add_subplot(5,1,2)

ax0.yaxis.set_major_formatter(yformatter)

line0=ax0.plot(x, val1_1d, label='Electron', linestyle = linestyles[0])
line0=ax0.plot(x, val2_1d, label=r'$D^+ ion$', linestyle = linestyles[1])


ymin = 0.0
ymax = 1.2 * max( val1_1d.max(), val2_1d.max() )

ax0.set_xlim((xmin, xmax))
ax0.set_ylabel("Pflux", fontsize = label_fontsize)
#ax0.yaxis.set_label_coords(ylabel_x, 0.5)


#ax0.grid(True)

ax0.annotate("(b)", xy=get_axis_limits(ax0), annotation_clip=False)

##============Eenergy flux======================================================
val1 = collect("/Diagnostic/", "heat_flux_left")
val1_1d = val1[:, 0, 0, 0]

val2 = collect("/Diagnostic/", "heat_flux_left")
val2_1d = val2[:, 0, 0, 1]

ax0=fig.add_subplot(5,1,3)

ax0.yaxis.set_major_formatter(yformatter)

line0=ax0.plot(x, val1_1d, label='Electron', linestyle = linestyles[0])
line0=ax0.plot(x, val2_1d, label=r'$D^+ ion$', linestyle = linestyles[1])


ymin = 0.0
ymax = 1.2 * max( val1_1d.max(), val2_1d.max() )

ax0.set_xlim((xmin, xmax))
#ax0.set_ylim((ymin, ymax))

ax0.set_ylabel("Hflux", fontsize = label_fontsize)
ax0.yaxis.set_label_coords(ylabel_x, 0.5)
#ax0.grid(True)
#ax0.legend(loc = 1)


ax0.annotate("(c)", xy=get_axis_limits(ax0), annotation_clip=False)


##============particle source======================================================
val1 = collect("/ParticleSource/", "load_dn")
val1_1d = val1[:, 0, 0, 0]
val2_1d = val1[:, 0, 0, 1]

ax0=fig.add_subplot(5,1,4)

line0=ax0.plot(x, val1_1d, label='Electron', linestyle = linestyles[0])
line0=ax0.plot(x, val2_1d, label=r'$D^+ ion$', linestyle = linestyles[1])


ax0.set_xlim((xmin, xmax))
#ax0.set_ylim((ymin, ymax))
ax0.set_xlabel(r"$t\ (\mu s)$", fontsize = label_fontsize)
ax0.set_ylabel("load_dn", fontsize = label_fontsize)
#ax0.grid(True)
#ax0.legend(loc = 1)


ax0.annotate("(d)", xy=get_axis_limits(ax0), annotation_clip=False)


##============source density======================================================
val1 = collect("/Fields/", "Rho_global_e_avg")
val1_1d = val1[:, 0, 0, :]
nx_half = int(val1_1d.shape[1] / 2)

ax0=fig.add_subplot(5,1,5)

line0=ax0.plot(x, val1_1d[:, nx_half], label='Electron', linestyle = linestyles[0])



ax0.set_xlim((xmin, xmax))
#ax0.set_ylim((ymin, ymax))
ax0.set_xlabel(r"$t\ (\mu s)$", fontsize = label_fontsize)
ax0.set_ylabel("source density", fontsize = label_fontsize)
#ax0.grid(True)
#ax0.legend(loc = 1)


ax0.annotate("(d)", xy=get_axis_limits(ax0), annotation_clip=False)


fig.savefig("flux.png")
#fig.show()       #when the program finishes,the figure disappears
#plt.axis('equal')
#plt.show()         #The command is OK
